/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { Arg } from './arg';
import { BasicBlock } from './basic_block';
import { Code } from './code';
import { InsertionSet } from './insertion_set';
import { Inst, nop, storeZero32 } from './inst';
import { Liveness } from './liveness';
import { StackSlot } from './stack_solt';
import { rangesOverlap, removeAllMatching, roundUpToMultipleOf, stackAlignmentBytes } from './util';
import { LocalCalc } from './liveness';

const AS_BYTESIZE_FORE: number = 4;
const AS_BYTESIZE_EIGHT: number = 8;
const AS_BYTESIZE_THIRTY_TWO: number = 32;
export function allocateStack(code: Code): void {
  if (code.frameSize > 0) {
    throw new Error('Frame size already determined');
  }
  let assignedEscapedStackSlots: StackSlot[] = new Array();
  let escapedStackSlotsWorklist: StackSlot[] = new Array();

  for (let slot of code.stackSlots) {
    if (slot.isLocked) {
      if (slot.offsetFromFP !== null) {
        assignedEscapedStackSlots.push(slot);
      } else {
        escapedStackSlotsWorklist.push(slot);
      }
    } else {
      if (slot.offsetFromFP !== null) {
        throw new Error('Offset already assigned');
      }
    }
  }
  while (escapedStackSlotsWorklist.length > 0) {
    let slot = escapedStackSlotsWorklist.pop();
    assign(slot, assignedEscapedStackSlots);
    assignedEscapedStackSlots.push(slot);
  }
  // Now we handle the spill slots.
  let liveness = new Liveness(StackSlot, code);
  let interference: Map<StackSlot, Set<StackSlot>> = new Map();

  for (let slot of code.stackSlots) {
    interference.set(slot, new Set<StackSlot>());
  }
  let slots: StackSlot[] = new Array();
  handerCodeWithLiveness(code, liveness, interference);
  for (let slot of code.stackSlots) {
    if (slot.offsetFromFP !== null) {
      continue;
    }
    assign(slot, assignedEscapedStackSlots.concat(Array.from(interference.get(slot))));
  }
  restCode(code);
}

function restCode(code: Code): void {
  let frameSizeForStackSlots = 0;
  for (let slot of code.stackSlots) {
    frameSizeForStackSlots = Math.max(frameSizeForStackSlots, -slot.offsetFromFP);
  }
  frameSizeForStackSlots = roundUpToMultipleOf(stackAlignmentBytes, frameSizeForStackSlots);

  for (let block of code) {
    for (let inst of block) {
      handerInstAndCode(inst, code);
    }
  }
  code.setFrameSize(frameSizeForStackSlots + code.callArgAreaSize);
  let insertionSet = new InsertionSet();
  for (let block of code) {
    for (let instIndex = 0; instIndex < block.size; instIndex++) {
      handerForInst(block, instIndex, insertionSet, code);
    }
    insertionSet.execute(block._insts);
  }
}

function handerForInst(block: BasicBlock, instIndex: number, insertionSet: InsertionSet, code: Code): void {
  let inst = block.at(instIndex);
  inst.forEachArg((arg, role, type, width) => {
    switch (arg.kind) {
      case Arg.stack:
        let slot = arg.stackSlot();
        if (Arg.isZDef(role) && slot.isSpill && slot.byteSize > Arg.bytes(width)) {
          if (slot.byteSize !== AS_BYTESIZE_EIGHT) {
            // print(`"Bad spill slot size for ZDef: ${slot.byteSize}}, width is ${width}"`);
          }
          if (width !== AS_BYTESIZE_THIRTY_TWO) {
            // print('Bad width for ZDef');
          }
          insertionSet.append(instIndex + 1, new Inst(storeZero32, [stackAddr(arg.offset + AS_BYTESIZE_FORE + slot.offsetFromFP, code.frameSize, width)]));
        }
        return stackAddr(arg.offset + slot.offsetFromFP, code.frameSize, width);
      case Arg.callArg:
        return stackAddr(arg.offset - code.frameSize, code.frameSize, width);
      default:
        return null;
    }
  });
}

function handerInstAndCode(inst: Inst, code: Code): void {
  for (let arg of inst.args) {
    if (arg.isCallArg) {
      if (arg.offset < 0) {
        throw new Error('Did not expect negative offset for callArg');
      }
      code.requestCallArgAreaSize(arg.offset + AS_BYTESIZE_EIGHT);
    }
  }
}

function handerCodeWithLiveness(code: Code, liveness: Liveness, interference: Map<StackSlot, Set<StackSlot>>): void {
  for (let block of code) {
    let localCalc = liveness.localCalc(block);
    let instIndex = block.size;
    while (instIndex > 0) {
      instIndex -= 1;
      let inst = block.at(instIndex);
      handerLocal(inst, localCalc);
      interfere(instIndex, block, interference, localCalc);
      localCalc.execute(instIndex);
    }
    interfere(-1, block, interference, localCalc);
    removeAllMatching(block._insts, (item) => {
      return item.opcode === nop;
    });
  }
}

function handerLocal(inst: Inst, localCalc: LocalCalc): void {
  if (!inst.hasNonArgEffects) {
    let ok = true; // 初始化OK为true
    inst.forEachArg((arg, role, type, width) => {
      if (Arg.isEarlyDef(role)) {
        ok = false;
        return null;
      }
      if (!Arg.isLateDef(role)) {
        return null;
      }
      if (!arg.isStack) {
        ok = false;
        return null;
      }
      let slot = arg.stackSlot();
      if (!slot.isSpill) {
        ok = false;
        return null;
      }
      if (localCalc.liveSet.has(slot)) {
        ok = false;
        return null;
      }
      return null;
    });
    if (ok) {
      inst.clear();
    }
  }
}

export function attemptAssignment(slot: StackSlot, offsetFromFP: number, otherSlots: StackSlot[]): boolean {
  let offsetFromFPPa = offsetFromFP;
  if (offsetFromFPPa > 0) {
    throw new Error('Expect negative offset');
  }

  offsetFromFPPa = -roundUpToMultipleOf(slot.alignment, -offsetFromFPPa);

  for (let otherSlot of otherSlots) {
    if (otherSlot.offsetFromFP === null) {
      continue;
    }
    let overlap = rangesOverlap(offsetFromFPPa, offsetFromFPPa + slot.byteSize, otherSlot.offsetFromFP, otherSlot.offsetFromFP + otherSlot.byteSize);
    if (overlap) {
      return false;
    }
  }
  slot.setOffsetFromFP(offsetFromFPPa);
  return true;
}

export function interfere(instIndex: number, block: BasicBlock, interference: Map<StackSlot, Set<StackSlot>>, localCalc: LocalCalc): void {
  Inst.forEachDef(StackSlot, block.get(instIndex), block.get(instIndex + 1), (slot, role, type, width) => {
    if (!slot.isSpill) {
      return null;
    }
    for (let otherSlot of localCalc.liveSet) {
      interference.get(slot).add(otherSlot);
      interference.get(otherSlot).add(slot);
    }
    return null;
  });
}

export function assign(slot: StackSlot, otherSlots: StackSlot[]): void {
  if (attemptAssignment(slot, -slot.byteSize, otherSlots)) {
    return;
  }
  for (let otherSlot of otherSlots) {
    if (otherSlot.offsetFromFP === null) {
      continue;
    }
    if (attemptAssignment(slot, otherSlot.offsetFromFP - slot.byteSize, otherSlots)) {
      return;
    }
  }
  throw new Error('Assignment failed');
}

export function stackAddr(offset: number, frameSize: number, width: number): Arg {
  return Arg.createStackAddr(offset, frameSize, width);
}
